package com.ds.lens.flink.monitor.common;

import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.apache.flink.streaming.connectors.elasticsearch.ElasticsearchSinkFunction;
import org.apache.flink.streaming.connectors.elasticsearch.util.RetryRejectedExecutionFailureHandler;
import org.apache.flink.streaming.connectors.elasticsearch6.ElasticsearchSink;
import org.apache.http.Header;
import org.apache.http.HttpHost;
import org.apache.http.message.BasicHeader;

import java.util.List;
import java.util.Map;

import static com.ds.lens.flink.monitor.MonitorJobConstants.ES_BULK_SIZE_CONFIG_FIELD;
import static com.ds.lens.flink.monitor.MonitorJobConstants.ES_INTERVAL_MS_CONFIG_FIELD;

/**
 * @author ethan
 */
@Slf4j
public class LensElasticsearchSink<T> extends ElasticsearchSink.Builder<T> {

    public LensElasticsearchSink(List<HttpHost> transportAddresses,
                                 ElasticsearchSinkFunction<T> elasticsearchSinkFunction,
                                 Map<String, String> esConfig) {
        super(transportAddresses, elasticsearchSinkFunction);

        if (StringUtils.isNotEmpty(esConfig.get(ES_BULK_SIZE_CONFIG_FIELD))) {
            setBulkFlushMaxActions(Integer.valueOf(esConfig.get(ES_BULK_SIZE_CONFIG_FIELD)));
        }
        if (StringUtils.isNotEmpty(esConfig.get(ES_INTERVAL_MS_CONFIG_FIELD))) {
            setBulkFlushInterval(Integer.valueOf(esConfig.get(ES_INTERVAL_MS_CONFIG_FIELD)));
        }
        Header[] headers = new Header[2];
        headers[0] = new BasicHeader("Connection", "Keep-Alive");
        headers[1] = new BasicHeader("Proxy-Connection", "Keep-Alive");
        setRestClientFactory(
                restClientFactory -> {
                    restClientFactory.setDefaultHeaders(headers);
                    restClientFactory.setMaxRetryTimeoutMillis(5 * 60 * 1000);
                    restClientFactory.setRequestConfigCallback(
                            requestConfigCallback -> {
                                requestConfigCallback.setConnectTimeout(5000);
                                requestConfigCallback.setSocketTimeout(40000);
                                requestConfigCallback.setConnectionRequestTimeout(1000);
                                return requestConfigCallback;
                            }
                    );
                }
        );
//        setBulkFlushBackoff(true);
//        setBulkFlushBackoffType(ElasticsearchSinkBase.FlushBackoffType.EXPONENTIAL);
//        setBulkFlushBackoffRetries(8);
//        setBulkFlushBackoffDelay(10000);
        setFailureHandler(new RetryRejectedExecutionFailureHandler());
    }

    public LensElasticsearchSink(List<HttpHost> transportAddresses,
                                 ElasticsearchSinkFunction<T> elasticsearchSinkFunction,
                                 Integer bulkFlushMaxActions, Long bulkFlushInterval) {
        super(transportAddresses, elasticsearchSinkFunction);

        if (null != bulkFlushMaxActions && bulkFlushMaxActions > 0) {
            setBulkFlushMaxActions(bulkFlushMaxActions);
        }
        if (null != bulkFlushInterval && bulkFlushInterval > 0) {
            setBulkFlushInterval(bulkFlushInterval);
        }
        Header[] headers = new Header[2];
        headers[0] = new BasicHeader("Connection", "Keep-Alive");
        headers[1] = new BasicHeader("Proxy-Connection", "Keep-Alive");
        setRestClientFactory(
                restClientFactory -> {
                    restClientFactory.setDefaultHeaders(headers);
                    restClientFactory.setMaxRetryTimeoutMillis(5 * 60 * 1000);
                    restClientFactory.setRequestConfigCallback(
                            requestConfigCallback -> {
                                requestConfigCallback.setConnectTimeout(5000);
                                requestConfigCallback.setSocketTimeout(40000);
                                requestConfigCallback.setConnectionRequestTimeout(1000);
                                return requestConfigCallback;
                            }
                    );
                }
        );

        setFailureHandler(new RetryRejectedExecutionFailureHandler());
    }
}
